package com.ds.infrastructure.hyperspace.container.flow.handler.security;

import com.google.common.util.concurrent.RateLimiter;
import com.ds.infrastructure.hyperspace.common.Container;
import com.ds.infrastructure.hyperspace.common.flow.FlowHandlerPipeline;
import com.ds.infrastructure.hyperspace.container.flow.context.HyperspaceDuplexFlowContext;
import com.ds.infrastructure.hyperspace.container.flow.handler.AbstractPreFlowHandler;
import com.ds.infrastructure.hyperspace.container.internal.operator.OperatorDefinition;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.StringUtils;

import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

import static com.ds.infrastructure.hyperspace.common.StateEnum.REQUEST_BLOCKED;

/**
 * usage eg:
 *  -   handler: rateLimiterPreHandler
 *      handlerParams: [value1|value2|valueN, qpsValue1|qpsValue2|qpsValueN, uri|ip|headerName, (operatorName,operatorArgs)...]
 * @author  Weishaoying
 */
@Slf4j
public class RateLimiterPreFlowHandler extends AbstractPreFlowHandler {

    private static final int INPUT_NAME_INDEX = 2;

    private Map<String, RateLimiter> valueRateLimiterMap = new ConcurrentHashMap<>();

    private volatile String inputName;

    private volatile List<OperatorDefinition> operatorDefinitionList;

    @Override
    public void handleInput(HyperspaceDuplexFlowContext duplexFlowContext, Object data) {
        if (data == duplexFlowContext.getRequestStartState()) {
            if (inputName == null) {
                return;
            }
            String inputValue = extractAndComputeInput(duplexFlowContext, inputName, operatorDefinitionList);
            if(StringUtils.isEmpty(inputValue)){
                return;
            }
            RateLimiter rateLimiter = valueRateLimiterMap.get(inputValue);
            if (rateLimiter == null || rateLimiter.tryAcquire()) {
                return;
            }
            duplexFlowContext.getEntryFlowOperation().interruptAndCloseContext(duplexFlowContext, REQUEST_BLOCKED,
                    "Rate limit was triggered, match item " + inputName +
                            ", match value " + inputValue + ", rate value + " + rateLimiter.getRate());
        }
    }

    @Override
    public String handlerName() {
        return "rateLimiterPreHandler";
    }

    @Override
    public void initHandlerParameters(Container container, FlowHandlerPipeline pipeline, String[] handlerParams) {
        if (handlerParams.length <= INPUT_NAME_INDEX) {
            return;
        }
        String[] valueArray = StringUtils.delimitedListToStringArray(
                handlerParams[0], OR_DELIMITER);
        String[] qpsArray = StringUtils.delimitedListToStringArray(
                handlerParams[1], OR_DELIMITER);
        if (valueArray.length != qpsArray.length) {
            log.warn("RateLimiterPreHandler init failed, args values or qps config invalid!");
            return;
        }
        if (valueArray.length > 0) {
            for (int i = 0; i < valueArray.length; i++) {
                valueRateLimiterMap.putIfAbsent(valueArray[i],
                        RateLimiter.create(Double.valueOf(qpsArray[i])));
            }
        }
        inputName = handlerParams[INPUT_NAME_INDEX];
        operatorDefinitionList = parseOperatorDefinition(container, handlerParams, INPUT_NAME_INDEX + 1);
    }

}
